import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def myplot2d(traj):
        # 绘制关节轨迹
        t=traj.t
        q=traj.q
        arrive=traj.arrive
        via=traj.via
        print(arrive)
        print(via)
        K, N = q.shape
        for i in range(N):
                plt.plot(t, q[:, i], label=f'Joint {i+1}')
        plt.scatter(arrive, via[:,0], marker='*', color='red', label='J1 Via')
        plt.scatter(arrive, via[:,1], marker='*', color='red', label='J2 Via')
        plt.scatter(arrive, via[:,2], marker='*', color='red', label='J3 Via')
        # 添加标签和标题
        plt.xlabel('Time')
        plt.ylabel('Joint Angle')
        plt.title('Joint Trajectory')
        plt.legend()
        plt.grid(True)
        plt.show()
def myplot3d(traj):
    # 创建3D图形对象
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # 绘制空间轨迹
    ax.plot(traj[:,0], traj[:,1], traj[:,2], label='Space Trajectory')

    # 添加标签和标题
    ax.set_xlabel('X-axis')
    ax.set_ylabel('Y-axis')
    ax.set_zlabel('Z-axis')
    ax.set_title('3D Space Trajectory')

    # 显示图例
    ax.legend()

    # 显示图形
    plt.show()
    
def linearPlanning(p1f, p2f, dotNum):
    '''
    *@brief:   笛卡尔空间直线规划函数
    *@date:    2021.9.12
    *@param:   p1f：起始点坐标，元组形式输入
    *@param:   p2f：终止点坐标，元组形式输入
    *@param:   dotNum：插值点数，以int形式输入
    *@returnParam:   pointf：规划后的点坐标list
    '''
    if not (isinstance(p1f, tuple) and isinstance(p2f, tuple)):
        print('两个路径点需要以元组的形式输入')
        return
    if len(p1f) != 3 or len(p2f) != 3:
        print('p1,p2中存在某点元组长度不为3，请检查')
        return
    deltax = (p2f[0] - p1f[0]) / dotNum
    deltay = (p2f[1] - p1f[1]) / dotNum
    deltaz = (p2f[2] - p1f[2]) / dotNum
    pointf = []  # 存储路径点
    currentIndex = 0
    while currentIndex <= dotNum:
        x = p1f[0] + deltax * currentIndex
        y = p1f[1] + deltay * currentIndex
        z = p1f[2] + deltaz * currentIndex
        pointf.append((x, y, z))
        currentIndex += 1
    return pointf


def showSpacePath(pointf, waypointf):
    '''
    *@brief:   显示空间轨迹的函数
    *@date:    2021.9.12
    *@param:   pointf：规划出来的空间点
    *@param:   waypointf：给定的路径点
    *@returnParam: 无，最后会显示图像
    '''
    xf = []
    yf = []
    zf = []
    xof = []
    yof = []
    zof = []
    color = []
    cMax = 1
    cMin = 0.9
    j = 0
    for i in pointf:
        xf.append(i[0])
        yf.append(i[1])
        zf.append(i[2])
        color.append(cMin + (cMax - cMin) * j / len(pointf))
        j += 1
    for i in waypointf:
        xof.append(i[0])
        yof.append(i[1])
        zof.append(i[2])
    fig = plt.figure()
    # plt.scatter(x, y)
    ax = Axes3D(fig, auto_add_to_figure=False)
    fig.add_axes(ax)
    ax.scatter(xf, yf, zf, c=color, alpha=0.3, cmap='summer', label='planned points')
    ax.scatter(xof, yof, zof, c='r', label='original points', s=70)
    ax.set_zlabel('Z', fontdict={'size': 15, 'color': 'black'})
    ax.set_ylabel('Y', fontdict={'size': 15, 'color': 'black'})
    ax.set_xlabel('X', fontdict={'size': 15, 'color': 'black'})
    ax.legend()
    plt.show()

if __name__ == "__main__":

    from roboticstoolbox import *
    import numpy as np
    puma = models.DH.Puma560()
    # # 下面是关节空间轨迹规划
    # via = np.array([[0,0,0,0,0,0],
    #                [np.pi/6,np.pi/5,np.pi/4,np.pi/4,np.pi/4,np.pi/4],
    #                [np.pi/2,np.pi/2,np.pi/2,np.pi/2,np.pi/2,np.pi/2],
    #                [np.pi*2/3,np.pi,np.pi,np.pi*2/3,np.pi*2/3,np.pi*2/3],
    #                ])
    # traj0 = mstraj(via, dt=0.5, tacc=3, qdmax=np.deg2rad(10))
    # T=puma.fkine(traj0.q)
    # myplot3d(T.t)
    # puma.plot(traj0.q, block=True)
    # myplot2d(traj0)
    # print(traj0.via)
    
    #下面是笛卡尔空间轨迹规划
    via = np.array([[0,0,400],
                    [20,0,380],
                    [40,40,350],
                    [80,40,300],
                    [80,80,300],
    ])
    traj0 = mstraj(via, dt=0.05, tacc=0.1, qdmax=0.01)
    traj0.plot(block=True)
    myplot3d(traj0.q)
    myplot2d(traj0)
    
    # #下面是笛卡尔空间直线规划
    # # 笛卡尔空间直线规划举例
    # p1 = (30, 0, 20)
    # p2 = (0, 40, 50)
    # dotNum = 50
    # # 以下为程序运行段，不用改动，只需要改动上面的两个点坐标和插入点数即可
    # pointss = linearPlanning(p1, p2, dotNum)
    # showSpacePath(pointss, [p1, p2])

    